import pickle  # 新增
import cv2
import numpy as np
import socket
import struct
import numpy as np
import threading  # 线程锁，防止数据竞争
import cv2
import time
import orjson

from queue import Queue, Empty
import zlib  # 添加压缩支持
def to_list_safe(x):
    if isinstance(x, np.ndarray):
        return x.tolist()
    elif isinstance(x, dict):
        return {k: to_list_safe(v) for k, v in x.items()}
    elif isinstance(x, list):
        return [to_list_safe(i) for i in x]
    return x
def get_local_ip():
    """获取当前设备的局域网 IP"""
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        # 连接到外部地址，但不发送数据，用于获取本地IP
        s.connect(('8.8.8.8', 80))  
        ip = s.getsockname()[0]
        s.close()
    except Exception:
        ip = '127.0.0.1'  # 失败时退回本地回环地址
    return ip
# 获取本地 IP 地址
local_ip = get_local_ip()
print(f"Local IP address: {local_ip}")
encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 90]

class StereoImagePublisher():
    def __init__(self, host=local_ip, port=8000):
        # 初始化 TCP 服务器
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        # 设置 TCP 发送/接收缓冲区大小
        BUFFER_SIZE = 2 * 100 * 1280  # 4MB
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, BUFFER_SIZE)
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, BUFFER_SIZE)
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)  # 允许端口复用
        self.server_socket.bind((host, port))
        self.server_socket.listen(2)  # 允许一个客户端连接
        print(f"Waiting for TCP connection on {host}:{port}...")

        # 监听客户端连接
        self.client_socket, self.client_address = self.server_socket.accept()
        #禁用Nagle算法（减少延迟）：
        self.client_socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        self.server_socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        print(f"Client connected: {self.client_address}")

        self.send_lock = threading.Lock()  # 添加发送锁
        self.compression_level = 0  # 压缩级别 (0-9)

    def send_data(self, data_type, data_bytes):
        if self.client_socket is None:
            print("No client connected. Skipping data send.")
            return
        try:
            # 压缩数据
            compressed_data = zlib.compress(data_bytes, self.compression_level)
            total_size = len(compressed_data)
            header = struct.pack("16sI", data_type.encode(), total_size)
            
            with self.send_lock:
                self.client_socket.sendall(header + compressed_data)
        except Exception as e:
            print(f"Send failed: {e}")
            self.client_socket.close()
            self.client_socket = None
            self.client_socket, self.client_address = self.server_socket.accept()
            print(f"Client reconnected: {self.client_address}")

    def __call__(self, ps_dict_2d, kps_2d_list, pts_3d_dict, pts_3d_list, left_image, right_image):
        try:
            # 批量处理数据
            data_pack = {
                "ps2d": to_list_safe(ps_dict_2d),
                "kps2d": to_list_safe(kps_2d_list),
                "pts3d": to_list_safe(pts_3d_dict),
                "ptlist3d": to_list_safe(pts_3d_list)
            }

            # 压缩并发送结构化数据
            packed_data = orjson.dumps(data_pack)
            self.send_data("packed_data", packed_data)

            # 验证图像有效性
            if left_image is None or right_image is None:
                print("警告: 接收到空图像，跳过图像发送")
                return

            if left_image.size == 0 or right_image.size == 0:
                print(f"警告: 图像尺寸无效 - left: {left_image.shape if hasattr(left_image, 'shape') else 'None'}, "
                      f"right: {right_image.shape if hasattr(right_image, 'shape') else 'None'}")
                return

            # 确保图像是有效的BGR格式
            if len(left_image.shape) == 2:  # 如果是灰度图，转为BGR
                left_image = cv2.cvtColor(left_image, cv2.COLOR_GRAY2BGR)
            if len(right_image.shape) == 2:  # 如果是灰度图，转为BGR
                right_image = cv2.cvtColor(right_image, cv2.COLOR_GRAY2BGR)

            # 压缩并发送图像
            try:
                _, left_encoded = cv2.imencode('.jpg', left_image, encode_param)
                _, right_encoded = cv2.imencode('.jpg', right_image, encode_param)

                # 合并图像数据
                combined_images = {
                    "left": left_encoded.tobytes(),
                    "right": right_encoded.tobytes()
                }
                packed_images = pickle.dumps(combined_images)
                self.send_data("images", packed_images)

            except cv2.error as cv_err:
                print(f"OpenCV编码错误: {cv_err}")
                print(f"图像信息 - left: dtype={left_image.dtype}, shape={left_image.shape}, "
                      f"right: dtype={right_image.dtype}, shape={right_image.shape}")
            except Exception as enc_err:
                print(f"图像编码异常: {enc_err}")

        except Exception as e:
            print(f"数据打包错误: {e}")
            import traceback
            traceback.print_exc()

    def close_sockets(self):
        if self.client_socket:
            self.client_socket.close()
        if self.server_socket:
            self.server_socket.close()
        print("All sockets closed.")
        
StereoImagePublisher_ins = StereoImagePublisher()   
class StereoDataSender:
    def __init__(self):
        self.data_queue = Queue(maxsize=10)  # 限制队列大小以避免内存占用过大
        self.running = True
        self.sender_thread = threading.Thread(target=self._sender_loop)
        self.last_send_time = time.time()
        self.min_send_interval = 0.033  # 约30fps
        self.sender_thread.start()
    
    def _sender_loop(self):
        while self.running:
            try:
                # current_time = time.time()
                # if current_time - self.last_send_time < self.min_send_interval:
                #     time.sleep(0.001)  # 短暂休眠以减少CPU使用
                #     continue
                    
                data = self.data_queue.get(timeout=0.01)
                if data is not None:
                    kps_dict_2d, kps_2d_list, pts_3d_dict, pts_3d_list, left_image, right_image = data
                    StereoImagePublisher_ins(kps_dict_2d, kps_2d_list, pts_3d_dict, pts_3d_list, left_image, right_image)
                    self.last_send_time = time.time()
            except Empty:
                continue
            except Exception as e:
                print(f"Error in sender thread: {e}")
                continue
    
    def send_data(self, kps_dict_2d, kps_2d_list, pts_3d_dict, pts_3d_list, left_image, right_image):
        try:
            # 如果队列满了，移除最旧的数据
            if self.data_queue.full():
                try:
                    self.data_queue.get_nowait()
                except Empty:
                    pass
            # 添加新数据
            self.data_queue.put((kps_dict_2d, kps_2d_list, pts_3d_dict, pts_3d_list, left_image, right_image))
        except Exception as e:
            print(f"Error queueing data: {e}")
    
    def stop(self):
        self.running = False
        if self.sender_thread.is_alive():
            self.sender_thread.join()

if __name__ == '__main__':
    stereo_image_publisher = StereoImagePublisher()
    while True:
        stereo_image_publisher.send_data()